package com.treasure.infrastructure.domain.llm.qianwen;

import com.alibaba.dashscope.embeddings.TextEmbeddingResultItem;
import com.alibaba.dashscope.exception.NoApiKeyException;
import com.treasure.domain.domain.EmbeddingVectorDomainService;
import com.treasure.domain.domain.entity.EmbeddingResultEntity;
import com.treasure.common.exception.ServiceException;
import com.treasure.common.types.Question;
import com.treasure.infrastructure.domain.llm.qianwen.client.QwenEmbeddingVectorClient;
import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.springframework.stereotype.Service;
import org.springframework.util.CollectionUtils;

import java.util.Collections;
import java.util.List;
import java.util.stream.Collectors;

/**
 * @author sunhongmin
 * @date 2023/8/27 10:02
 * @description
 */
@Service
@RequiredArgsConstructor
@Slf4j
public class EmbeddingVectorDomainServiceImpl implements EmbeddingVectorDomainService {

    private final QwenEmbeddingVectorClient qwenEmbeddingVectorClient;

    @Override
    public List<EmbeddingResultEntity> textEmbedding(List<Question> questionList) {
        if (CollectionUtils.isEmpty(questionList)) {
            return Collections.emptyList();
        }

        List<String> questions = questionList.stream().map(Question::getValue).collect(Collectors.toList());
        List<TextEmbeddingResultItem> textEmbeddingResultItems;
        try {
            textEmbeddingResultItems = qwenEmbeddingVectorClient.textEmbeddingVector(questions);
        } catch (NoApiKeyException e) {
            log.error("textEmbedding NoApiKeyException e:", e);
            throw new ServiceException(e.getMessage());
        }

        if (CollectionUtils.isEmpty(textEmbeddingResultItems)) {
            return Collections.emptyList();
        }

        return textEmbeddingResultItems.stream().map(it -> {
            EmbeddingResultEntity embeddingResultEntity = new EmbeddingResultEntity();
            embeddingResultEntity.setTextIndex(it.getTextIndex());
            embeddingResultEntity.setEmbeddings(it.getEmbedding());
            return embeddingResultEntity;
        }).collect(Collectors.toList());
    }
}
